import os
import numpy
import pandas as pd
import time
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
def savepic(names):
    file_tmp = open("./{}".format(names), 'r')
    tmp_lines = [i for i in file_tmp]
    file_tmp.close()
    aa_number = []
    plddt = []
    tmp_lines_clean = []

    for i in tmp_lines:
        if i[0:4] == "ATOM":
            tmp_lines_clean.append(i)

    file_tmp = open("./{}_clean.pdb".format(names), 'w+')
    file_tmp.writelines(tmp_lines_clean)
    file_tmp.close()

    file_tmp = open("./{}_clean.pdb".format(names), 'r')
    tmp_lines = [i for i in file_tmp]
    file_tmp.close()

    for i in range(len(tmp_lines)):
        if int(tmp_lines[i][22:30]) not in aa_number:
            aa_number.append(int(tmp_lines[i][22:30]))
            plddt.append(float(tmp_lines[i][61:66]))
    dict = {'aa_number': aa_number, 'plddt': plddt}

    df = pd.DataFrame(dict)
    print(df)



    x = df['aa_number']
    y1 = df['plddt']

    # Plot Line1 (Left Y Axis)
    fig, ax1 = plt.subplots(1, 1, figsize=(16, 9), dpi=600)
    ax1.plot(x, y1, color='tab:red')

    # Decorations
    # ax1 (left Y axis)
    ax1.set_xlabel('Amino acid number', fontsize=24)
    # ax1.tick_params(axis='x', rotation=0, labelsize=10,labelcolor='tab:red')
    ax1.tick_params(axis='x', rotation=0, labelsize=10)
    # 显示所有横坐标
    # plt.xticks(x)
    ax1.xaxis.set_major_locator(ticker.MultipleLocator(base=15))

    ax1.set_ylabel('Plddt(0~100)', fontsize=24)
    # ax1.tick_params(axis='y', rotation=0, labelcolor='tab:red' )
    ax1.tick_params(axis='y', rotation=0)
    ax1.grid(alpha=.2)

    fig.tight_layout()
    plt.savefig('./{}.png'.format(names))




list = ["ranked_0.pdb","ranked_1.pdb","ranked_2.pdb","ranked_3.pdb","ranked_4.pdb"]

for n in list:
    savepic(n)

for i in range(5):
    if os.path.exists("./ranked_{}.pdb_clean.pdb".format(i)):
        os.remove("./ranked_{}.pdb_clean.pdb".format(i))
    else:
        print("The file ranked_{}.pdb_clean.pdb does not exist".format(i))
